import asyncio
from contextlib import asynccontextmanager


# 模拟一个异步连接池
class AsyncConnectionPool:
    def __init__(self, size):
        self.size = size
        self.connections = []

    async def warmup(self):
        """预热连接池"""
        print("Warming up connection pool...")
        for i in range(self.size):
            # 模拟异步创建连接
            self.connections.append(f"Connection-{i}")
        print("Pool is ready!")

    async def close(self):
        """关闭所有连接"""
        print("Closing all connections...")
        self.connections.clear()
        print("Pool is closed.")

    async def get_connection(self):
        """从池中获取一个连接（模拟）"""
        if not self.connections:
            raise Exception("No available connections!")
        # 模拟异步获取（例如等待一个可用的信号量）
        await asyncio.sleep(0.1)
        return self.connections.pop()

    async def return_connection(self, conn):
        """将连接归还给池（模拟）"""
        # 模拟异步归还
        await asyncio.sleep(0.1)
        self.connections.append(conn)


# 创建全局连接池
pool = AsyncConnectionPool(2)


@asynccontextmanager
async def get_db_connection():
    """
    一个异步上下文管理器，用于从全局池中获取和归还数据库连接。
    """
    conn = None
    try:
        # 1. __aenter__: 获取连接
        conn = await pool.get_connection()
        print(f"Acquired connection: {conn}")
        # 2. yield: 将连接提供给 with 块使用
        yield conn
        # 4. 如果 with 块成功，归还连接
        print("Operation succeeded, returning connection.")
        await pool.return_connection(conn)
    except Exception as e:
        # 5. 如果 with 块失败，也要归还连接，但可能记录日志等
        print(f"Operation failed: {e}. Returning connection.")
        if conn:
            await pool.return_connection(conn)
        # 重新抛出异常，不抑制它
        raise
    # finally 块在这里被 asynccontextmanager 隐式处理了


async def perform_database_operation(query):
    """
    使用上面定义的上下文管理器来执行数据库操作
    """
    async with get_db_connection() as connection:
        # 在这里使用连接执行操作
        print(f"Executing query '{query}' with {connection}")
        # 模拟一个可能成功也可能失败的异步操作
        await asyncio.sleep(0.5)
        if "fail" in query:
            raise RuntimeError("Database operation failed!")
        return "SUCCESS"


async def main():
    # 首先预热连接池
    await pool.warmup()

    # 执行一些成功的操作
    try:
        result = await perform_database_operation("SELECT * FROM users")
        print(f"Result: {result}")
    except Exception as e:
        print(f"Caught error: {e}")

    print("----")

    # 执行一个会失败的操作
    try:
        result = await perform_database_operation("INSERT INTO users fail")
        print(f"Result: {result}")
    except Exception as e:
        print(f"Caught error: {e}")

    # 最后关闭连接池
    await pool.close()


asyncio.run(main())
